package com.mjlf.algorithms.search;

import org.junit.Test;

/**
 * @auther a123
 * @create 2018-09-17 15:33
 * @desc 红黑树实现
 */
public class RedBlackTree<Key extends Comparable<Key>, Value> {

    public Node root;

    enum RedBlack {
        RED,
        BLACK
    }

    class Node {
        private Key key;
        private Value value;
        private Node left, right;
        private int numNode;
        private RedBlack color;

        public Node(Key key, Value value, int numNode) {
            this.key = key;
            this.value = value;
            this.numNode = numNode;
        }

        @Override
        public String toString() {
            String res = "{key:" + key + " color:" + color + "}";
            return res;
        }
    }

    public int size() {
        return size(root);
    }

    private int size(Node node) {
        if (node == null) {
            return 0;
        }
        return node.numNode;
    }

    /**
     * 插入数据
     *
     * @param key
     * @param value
     */
    public void put(Key key, Value value) {
        root = put(root, key, value);
        root.color = RedBlack.BLACK;
    }

    private Node put(Node node, Key key, Value value) {
        if (node == null) {
            node = new Node(key, value, 1);
            node.color = RedBlack.RED;
        }
        int cmp = key.compareTo(node.key);
        if (cmp < 0) {
            node.left = put(node.left, key, value);
        } else if (cmp > 0) {
            node.right = put(node.right, key, value);
        } else {
            node.value = value;
        }
        //进行左旋
        if (node.right != null) {
            if (node.right.color == RedBlack.RED) {
                node = leftLateral(node);
            }
        }
        if (node.left != null && node.left.left != null) {
            if (node.left.color == RedBlack.RED && node.left.left.color == RedBlack.RED) {
                //右旋
                node = rightLateral(node);
            }
        }
        if (node.left != null && node.right != null) {
            if (node.left.color == RedBlack.RED && node.right.color == RedBlack.RED) {
                node = change(node);
            }
        }
        return node;
    }

    /**
     * 如果一个节点的右子节点是红色，进行左旋
     *
     * @param node
     */
    private Node leftLateral(Node node) {
        Node x = node.right;
        node.right = x.left;
        x.left = node;
        x.color = node.color;
        node.color = RedBlack.RED;
        return x;
    }

    @Test
    public void test() {
        RedBlackTree<Integer, Integer> tree = new RedBlackTree<>();
        tree.put(12, 9);
        tree.put(1, 7);
        tree.put(9, 8);
        tree.put(2, 8);
        tree.put(0, 8);
        tree.put(11, 8);
        tree.put(7, 8);
        tree.put(19, 8);
        tree.put(4, 8);
        tree.put(15, 8);
        tree.put(18, 8);
        tree.put(5, 8);
        tree.put(14, 8);
        tree.put(13, 8);
        tree.put(16, 8);
        tree.put(6, 8);
        tree.put(3, 8);
        tree.put(8, 8);
        tree.put(17, 8);

        tree.delete(9);

        tree.show();
    }

    /**
     * 右旋
     * 当出现两项两个左红节点的时候
     *
     * @param upNode
     */
    private Node rightLateral(Node upNode) {
        Node node = upNode.left;

        upNode.left = node.right;
        node.right = upNode;
        node.left.color = RedBlack.BLACK;
        node.right.color = RedBlack.BLACK;
        node.color = RedBlack.RED;
        return node;
    }

    /**
     * 一个节点有两个红子节点， 进行转换， 将父节点转换为红节点
     *
     * @param mid
     */
    private Node change(Node mid) {
        mid.left.color = RedBlack.BLACK;
        mid.right.color = RedBlack.BLACK;
        mid.color = RedBlack.RED;
        return mid;
    }

    public void delete(Key key) {
        root = delete(root, key);
    }

    private Node delete(Node node, Key key) {
        if (node == null) {
            return null;
        }
        int cmp = key.compareTo(node.key);
        if (cmp == 0) {//判断相等， 删除
            Node beforeOrNext = foreRunner(node.left) != null ? foreRunner(node.left) : succeed(node.right);
            if (beforeOrNext == null) {//没有前驱或后继， 表示它是叶子节点， 直接删除
                return null;
            } else {
                //复制前驱或后继节点数据到当前需要删除的节点
                node.key = beforeOrNext.key;
                node.value = beforeOrNext.value;
                Node dNode = deleteBeforeOrEnd(beforeOrNext);
                if(dNode != null){
                    beforeOrNext.key = dNode.key;
                    beforeOrNext.value = dNode.value;
                    beforeOrNext.left = dNode.left;
                    beforeOrNext.right = dNode.right;
                }else {
                    beforeOrNext.color = node.color;
                    return beforeOrNext;
                }
            }
        } else if (cmp < 0) {
            node.left = delete(node.left, key);
        } else {
            node.right = delete(node.right, key);
        }
        return node;
    }

    private Node deleteBeforeOrEnd(Node beforeOrNext) {
        //删除前驱或后继节点
        if (beforeOrNext.color == RedBlack.RED) {//前驱或后续为红色节点
            if (beforeOrNext.left != null) {
                beforeOrNext.left.color = beforeOrNext.color;
                beforeOrNext = beforeOrNext.left;
            } else if (beforeOrNext.right != null) {
                beforeOrNext.right.color = beforeOrNext.color;
                beforeOrNext = beforeOrNext.right;
            } else {
                beforeOrNext = null;
            }
        } else {
            if (beforeOrNext.left != null) {//前驱或后续节点为黑色节点
                beforeOrNext = beforeOrNext.left;
                beforeOrNext.color = RedBlack.RED;
            } else if (beforeOrNext.right != null) {
                beforeOrNext = beforeOrNext.right;
                beforeOrNext.color = RedBlack.RED;
            } else {
                beforeOrNext = null;
            }
        }
        return beforeOrNext;
    }

    /**
     * 查找前驱
     *
     * @param node
     *
     * @return
     */
    public Node foreRunner(Node node) {
        if (node == null) {
            return null;
        }
        if (node.right != null) {
            return foreRunner(node.right);
        } else {
            return node;
        }
    }

    /**
     * 查找后继
     *
     * @param node
     *
     * @return
     */
    public Node succeed(Node node) {
        if (node == null) {
            return null;
        }
        if (node.left != null) {
            return succeed(node.left);
        } else {
            return node;
        }
    }

    public void show() {
        show(root);
        System.out.println();
    }

    private void show(Node node) {
        if (node != null) {
            System.out.print(node);
            if (node.left != null) {
                show(node.left);
            }
            if (node.right != null) {
                show(node.right);
            }
        }
    }
}
